# ------------------------------------------------------------------------------------
# Sparse DETR
# Copyright (c) 2021 KakaoBrain. All Rights Reserved.
# Licensed under the Apache License, Version 2.0 [see LICENSE for details]
# ------------------------------------------------------------------------------------
# Modified from Deformable DETR (https://github.com/fundamentalvision/Deformable-DETR)
# Copyright (c) 2020 SenseTime. All Rights Reserved.
# ------------------------------------------------------------------------------------
# Modified from DETR (https://github.com/facebookresearch/detr)
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
# ------------------------------------------------------------------------------------


"""
Deformable DETR model and criterion classes.
"""
import torch
import torch.nn.functional as F
from torch import nn
import math

from util.misc import (NestedTensor, nested_tensor_from_tensor_list,
                        inverse_sigmoid)
from .module.modules import MLP

from .backbones.sparse_backbone import build_backbone
from .necks.sparse_transformer import build_transformer
from .matcher.matcher import build_matcher
from .losses.setcriterion import SetCriterion
from .postprocess.postprocess import PostProcess

from .necks.module.utils import _get_clones

def build(args):
    num_classes = 20 if args.dataset_file != 'coco' else 91
    if args.dataset_file == "coco_panoptic":
        num_classes = 250
    device = torch.device(args.device)

    backbone = build_backbone(args)
    transformer = build_transformer(args)
    
    model = SparseDETR(
        backbone,
        transformer,
        num_classes=num_classes,
        num_queries=args.num_queries,
        num_feature_levels=args.num_feature_levels,
        aux_loss=args.aux_loss,
        with_box_refine=args.with_box_refine,
        two_stage=args.two_stage,
        args=args,
    )
    # if args.masks:
    #     model = DETRsegm(model, freeze_detr=(args.frozen_weights is not None))
    matcher = build_matcher(args, 'focal_loss')
    weight_dict = get_weight_dict(args)
    losses = get_losses(args)

    # num_classes, matcher, weight_dict, losses, focal_alpha=0.25
    criterion = SetCriterion(args, matcher, weight_dict, losses, focal_alpha=args.focal_alpha, 
                             loss_type='focal_loss')
    criterion.to(device)
    postprocessors = {'bbox': PostProcess(args, trans_type='ddetr')}
    # if args.masks:
    #     postprocessors['segm'] = PostProcessSegm()
    #     if args.dataset_file == "coco_panoptic":
    #         is_thing_map = {i: i <= 90 for i in range(201)}
    #         postprocessors["panoptic"] = PostProcessPanoptic(is_thing_map, threshold=0.85)

    return model, criterion, postprocessors



class SparseDETR(nn.Module):
    """ This is the Deformable DETR module that performs object detection """
    def __init__(self, backbone, transformer, num_classes, num_queries, num_feature_levels,
                 aux_loss=True, with_box_refine=False, two_stage=False, args=None, num_verb_classes=0):
        """ Initializes the model.
        Parameters:
            backbone: torch module of the backbone to be used. See backbone.py
            transformer: torch module of the transformer architecture. See transformer.py
            num_classes: number of object classes
            num_queries: number of object queries, ie detection slot. This is the maximal number of objects
                         DETR can detect in a single image. For COCO, we recommend 100 queries.
            aux_loss: True if auxiliary decoding losses (loss at each decoder layer) are to be used.
            with_box_refine: iterative bounding box refinement
            two_stage: two-stage Deformable DETR
        """
        super().__init__()

        self.backbone = backbone
        self.aux_loss = aux_loss
        self.with_box_refine = with_box_refine
        self.two_stage = two_stage

        self.use_enc_aux_loss = args.use_enc_aux_loss
        self.rho = args.rho

        self.num_queries = num_queries
        self.transformer = transformer
        hidden_dim = transformer.d_model
        self.hidden_dim = hidden_dim

        self.query_embeds = None
        if not self.two_stage:
            self.query_embed = nn.Embedding(num_queries, hidden_dim * 2)

        # will be splited into query_embed(query_pos) & tgt later


        self.num_feature_levels = num_feature_levels
        if num_feature_levels > 1:
            num_backbone_outs = len(backbone.strides)
            input_proj_list = []
            for _ in range(num_backbone_outs):
                in_channels = backbone.num_channels[_]
                input_proj_list.append(nn.Sequential(
                    nn.Conv2d(in_channels, hidden_dim, kernel_size=1),
                    nn.GroupNorm(32, hidden_dim),
                ))
            for _ in range(num_feature_levels - num_backbone_outs):
                input_proj_list.append(nn.Sequential(
                    nn.Conv2d(in_channels, hidden_dim, kernel_size=3, stride=2, padding=1),
                    nn.GroupNorm(32, hidden_dim),
                ))
                in_channels = hidden_dim
            self.input_proj = nn.ModuleList(input_proj_list)
        else:
            self.input_proj = nn.ModuleList([
                nn.Sequential(
                    nn.Conv2d(backbone.num_channels[0], hidden_dim, kernel_size=1),
                    nn.GroupNorm(32, hidden_dim),
                )])
        for proj in self.input_proj:
            nn.init.xavier_uniform_(proj[0].weight, gain=1)
            nn.init.constant_(proj[0].bias, 0)
            
        self.build_head(num_classes, num_verb_classes)

    def build_head(self, num_classes, num_verb_classes):
        hidden_dim = self.hidden_dim
        self.class_embed = nn.Linear(hidden_dim, num_classes)
        self.bbox_embed = MLP(hidden_dim, hidden_dim, output_dim=4, num_layers=3)
        
        prior_prob = 0.01
        bias_value = -math.log((1 - prior_prob) / prior_prob)
        self.class_embed.bias.data = torch.ones(num_classes) * bias_value
        nn.init.constant_(self.bbox_embed.layers[-1].weight.data, 0)
        nn.init.constant_(self.bbox_embed.layers[-1].bias.data, 0)


        # hack implementation: a list of embedding heads (see the order)
        # n: dec_layers / m: enc_layers
        # [dec_0, dec_1, ..., dec_n-1, encoder, backbone, enc_0, enc_1, ..., enc_m-2]
        
        # at each layer of decoder (by default)
        num_pred = self.transformer.decoder.num_layers
        if self.two_stage:
            # at the end of encoder
            num_pred += 1  
        if self.use_enc_aux_loss:
            # at each layer of encoder (excl. the last)
            num_pred += self.transformer.encoder.num_layers - 1 
             
        if self.with_box_refine or self.use_enc_aux_loss:
            # individual heads with the same initialization
            self.class_embed = _get_clones(self.class_embed, num_pred)
            self.bbox_embed = _get_clones(self.bbox_embed, num_pred)
            nn.init.constant_(self.bbox_embed[0].layers[-1].bias.data[2:], -2.0)
        else:
            # shared heads
            nn.init.constant_(self.bbox_embed.layers[-1].bias.data[2:], -2.0)
            self.class_embed = nn.ModuleList([self.class_embed for _ in range(num_pred)])
            self.bbox_embed = nn.ModuleList([self.bbox_embed for _ in range(num_pred)])
            
        if self.two_stage:
            # hack implementation
            self.transformer.decoder.class_embed = self.class_embed
            self.transformer.decoder.bbox_embed = self.bbox_embed            
            for box_embed in self.transformer.decoder.bbox_embed:
                nn.init.constant_(box_embed.layers[-1].bias.data[2:], 0.0)
                
        if self.use_enc_aux_loss:
            # the output from the last layer should be specially treated as an input of decoder
            num_layers_excluding_the_last = self.transformer.encoder.num_layers - 1
            self.transformer.encoder.aux_heads = True
            self.transformer.encoder.class_embed = self.class_embed[-num_layers_excluding_the_last:]
            self.transformer.encoder.bbox_embed = self.bbox_embed[-num_layers_excluding_the_last:] 
            for box_embed in self.transformer.encoder.bbox_embed:
                nn.init.constant_(box_embed.layers[-1].bias.data[2:], 0.0)
 

    def forward(self, samples: NestedTensor, targets=None):
        srcs, masks, pos = self.backbone_forward(samples)
        if not self.two_stage:
            outputs = self.neck_forward(srcs, masks, pos, self.query_embed.weight)
        else:
            outputs = self.neck_forward(srcs, masks, pos, None)
        out = self.head_forward(outputs)
        out.update(outputs)
        return out


    def backbone_forward(self, samples: NestedTensor):
        if not isinstance(samples, NestedTensor):
            samples = nested_tensor_from_tensor_list(samples)
        features, pos = self.backbone(samples)

        srcs, masks = [], []
        for l, feat in enumerate(features):
            src, mask = feat.decompose()
            src = self.input_proj[l](src)
            srcs.append(src)
            masks.append(mask)
            assert mask is not None

        #   假设特征要更高维度，就最后一层映射
        if self.num_feature_levels > len(srcs):
            _len_srcs = len(srcs)
            for l in range(_len_srcs, self.num_feature_levels):
                if l == _len_srcs:
                    src = features[-1].tensors
                    src = self.input_proj[l](src)
                else:
                    src = srcs[-1]
                    src = self.input_proj[l](src)
                m = samples.mask
                mask = F.interpolate(m[None].float(), size=src.shape[-2:]).to(torch.bool)[0]
                pos_l = self.backbone[1](NestedTensor(src, mask)).to(src.dtype)
                srcs.append(src)
                masks.append(mask)
                pos.append(pos_l)

        return srcs, masks, pos


    def neck_forward(self, srcs, masks, pos, query_embeds):
        outputs = self.transformer(srcs, masks, pos, query_embeds)
        return outputs


    def head_forward(self, outputs):
        #   backbone
        backbone_topk_proposals = outputs['backbone_topk_proposals']
        backbone_mask_prediction = outputs['backbone_mask_prediction']
        #   ddetr
        hs = outputs['out_query']
        init_reference = outputs['init_reference']
        inter_references = outputs['inter_references']
        #   two-stage
        enc_outputs_class = outputs['enc_outputs_class']
        enc_outputs_coord_unact = outputs['enc_outputs_coord_unact']
        #   use_enc_aux_loss
        enc_inter_outputs_class = outputs['enc_inter_outputs_class']
        enc_inter_outputs_coord = outputs['enc_inter_outputs_coord']
                
        sampling_locations_enc = outputs['sampling_locations_enc']
        attn_weights_enc = outputs['attn_weights_enc']
        sampling_locations_dec = outputs['sampling_locations_dec']
        attn_weights_dec = outputs['attn_weights_dec']
        spatial_shapes = outputs['spatial_shapes']
        level_start_index = outputs['level_start_index']  
        masks = outputs['masks']
        

        outputs_classes = []
        outputs_coords = []
        for lvl in range(len(hs)):
            # lvl: level of decoding layer
            outputs_class = self.class_embed[lvl](hs[lvl])
            outputs_classes.append(outputs_class)
            
            outputs_coord = self.bbox_embed[lvl](hs[lvl])
            assert init_reference is not None and inter_references is not None
            if lvl == 0:
                reference = init_reference
            else:
                reference = inter_references[lvl - 1]
            reference = inverse_sigmoid(reference)
            if reference.shape[-1] == 4:
                outputs_coord += reference
            else:
                assert reference.shape[-1] == 2
                outputs_coord[..., :2] += reference
            outputs_coord = outputs_coord.sigmoid()
            outputs_coords.append(outputs_coord)
            
        outputs_class = torch.stack(outputs_classes)
        outputs_coord = torch.stack(outputs_coords)

        # the topmost layer output
        out = {
            "pred_logits": outputs_class[-1],
            "pred_boxes": outputs_coord[-1],
            "sampling_locations_enc": sampling_locations_enc,
            "attn_weights_enc": attn_weights_enc,
            "sampling_locations_dec": sampling_locations_dec,
            "attn_weights_dec": attn_weights_dec,
            "spatial_shapes": spatial_shapes,
            "level_start_index": level_start_index,
        }
        if self.aux_loss and self.training:
            # make loss from every intermediate layers (excluding the last one)
            out['aux_outputs'] = self._set_aux_loss(outputs_class[:-1], outputs_coord[:-1])

        if backbone_topk_proposals is not None:
            out["backbone_topk_proposals"] = backbone_topk_proposals

        if self.rho:
            out["backbone_mask_prediction"] = backbone_mask_prediction
            out["sparse_token_nums"] = self.transformer.sparse_token_nums
            
        if self.two_stage:
            enc_outputs_coord = enc_outputs_coord_unact.sigmoid()
            out['enc_outputs'] = {'pred_logits': enc_outputs_class, 
                                  'pred_boxes': enc_outputs_coord}

        if self.use_enc_aux_loss:
            out['aux_outputs_enc'] = self._set_aux_loss(enc_inter_outputs_class, 
                                                        enc_inter_outputs_coord)


        out['mask_flatten'] = torch.cat([m.flatten(1) for m in masks], 1)

        return out


    @torch.jit.unused
    def _set_aux_loss(self, outputs_class, outputs_coord):
        # this is a workaround to make torchscript happy, as torchscript
        # doesn't support dictionary with non-homogeneous values, such
        # as a dict having both a Tensor and a list.
        return [{'pred_logits': a, 'pred_boxes': b}
                for a, b in zip(outputs_class, outputs_coord)]


def get_weight_dict(args):
    weight_dict = {'loss_ce': args.cls_loss_coef, 'loss_bbox': args.bbox_loss_coef}
    weight_dict['loss_giou'] = args.giou_loss_coef
    
    weight_dict['loss_mask_prediction'] = args.mask_prediction_coef
    if args.masks:
        weight_dict["loss_mask"] = args.mask_loss_coef
        weight_dict["loss_dice"] = args.dice_loss_coef
        
    # TODO this is a hack
    aux_weight_dict = {}
    
    if args.aux_loss:
        for i in range(args.dec_layers - 1):
            aux_weight_dict.update({k + f'_{i}': v for k, v in weight_dict.items()})
            
    if args.two_stage:
        aux_weight_dict.update({k + f'_enc': v for k, v in weight_dict.items()})
        
    if args.use_enc_aux_loss:
        for i in range(args.enc_layers - 1):
            aux_weight_dict.update({k + f'_enc_{i}': v for k, v in weight_dict.items()})
            
    if args.rho:
        aux_weight_dict.update({k + f'_backbone': v for k, v in weight_dict.items()})
        
    if aux_weight_dict:
        weight_dict.update(aux_weight_dict)
    return weight_dict

def get_losses(args):
    losses = ['labels', 'boxes', 'cardinality', "corr"]
    if args.masks:
        losses += ["masks"]
    if args.rho:
        losses += ["mask_prediction"]
    return losses


